using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace Decision_Tree
{
	/* Decision Tree */
	public class DecisionTree 
	{
		DecisionTreeNode root;
		ArrayList examples = new ArrayList();
		ArrayList attributeList = new ArrayList();
		
		/* Public interface */
		public DecisionTree(bool bKnowledgeEnabled)
		{
			/* Sets is a dictionary of lists of examples that have a given value */			
			/*
			 * Complete Example set
			 * bCloseToGoal, bCloseToFood, bCloseToHerder, bTimeLeft, bHasFood
			 */
			
			if(!bKnowledgeEnabled)
			{
				//Dumb as
				print("Using Limited Knowledge");
				examples.Add(new Example(true, true, true, true, true,  Actions.Lost));
			}
			else
			{
				//Super smart like tom :)
				print("Using Full Knowledge");
				examples.Add(new Example(true, true, true, true, true,  Actions.EvadeHerder));
				examples.Add(new Example(true, true, true, true, false,  Actions.EvadeHerder));
				examples.Add(new Example(true, true, true, false, true,  Actions.EvadeHerder));
				examples.Add(new Example(true, true, true, false, false,  Actions.EvadeHerder));
				examples.Add(new Example(true, true, false, true, true,  Actions.MoveTowardsGoal)); //Works
				examples.Add(new Example(true, true, false, true, false,  Actions.MoveTowardsFood));
				examples.Add(new Example(true, true, false, false, true,  Actions.MoveTowardsGoal));
				examples.Add(new Example(true, true, false, false, false,  Actions.MoveTowardsFood));
				examples.Add(new Example(true, false, true, true, true,  Actions.EvadeHerder));
				examples.Add(new Example(true, false, true, true, false,  Actions.EvadeHerder));
				examples.Add(new Example(true, false, true, false, true,  Actions.EvadeHerder));
				examples.Add(new Example(true, false, true, false, false,  Actions.EvadeHerder));
				examples.Add(new Example(true, false, false, true, true,  Actions.MoveTowardsGoal));
				examples.Add(new Example(true, false, false, true, false,  Actions.MoveTowardsGoal));
				examples.Add(new Example(true, false, false, false, true,  Actions.MoveTowardsGoal));
				examples.Add(new Example(true, false, false, false, false,  Actions.MoveTowardsGoal));
				examples.Add(new Example(false, true, true, true, true,  Actions.EvadeHerder));
				examples.Add(new Example(false, true, true, true, false,  Actions.EvadeHerder)); //Works
				examples.Add(new Example(false, true, true, false, true,  Actions.EvadeHerder));
				examples.Add(new Example(false, true, true, false, false,  Actions.EvadeHerder));
				examples.Add(new Example(false, true, false, true, true,  Actions.MoveTowardsFood));
				examples.Add(new Example(false, true, false, true, false,  Actions.MoveTowardsFood));
				examples.Add(new Example(false, true, false, false, true,  Actions.MoveTowardsGoal));
				examples.Add(new Example(false, true, false, false, false,  Actions.MoveTowardsFood));
				examples.Add(new Example(false, false, true, true, true,  Actions.EvadeHerder));
				examples.Add(new Example(false, false, true, true, false,  Actions.EvadeHerder));
				examples.Add(new Example(false, false, true, false, true,  Actions.EvadeHerder)); // T
				examples.Add(new Example(false, false, true, false, false,  Actions.EvadeHerder));
				examples.Add(new Example(false, false, false, true, true,  Actions.MoveTowardsFood));
				examples.Add(new Example(false, false, false, true, false,  Actions.MoveTowardsGoal));
				examples.Add(new Example(false, false, false, false, true,  Actions.MoveTowardsGoal));
				examples.Add(new Example(false, false, false, false, false,  Actions.MoveTowardsGoal));
			}
			/*
			 * End Complete example set
			 */
			
			/* Sets up array list of all the attributes */	
			attributeList.Add(Attribute.bCloseToFood);
			attributeList.Add(Attribute.bCloseToGoal);
			attributeList.Add(Attribute.bCloseToHerder);
			attributeList.Add(Attribute.bHasFood);
			attributeList.Add(Attribute.bTimeLeft);
			
			/* begin tree construction */
			
			//handel the root case based on entropy outside the function - because C# is pass by value?
			if(this.entropy(examples) <= 0)
			{
				Console.WriteLine("EARLY CUT");
				
				//this is only one possible outcome so we dont need to start the recursion
				root = new Action();
				root.action = ((Example)examples[0]).action;
				return;
			}
			else
			{
				Console.WriteLine("START RECUSION");
				root = new Decision();
				this._makeTree(examples,attributeList,root);
			}
		}
		
		private void print(string asdf)
		{
			MonoBehaviour.print(asdf);	
		}
		
		public Actions makeDecision(Knowledge agent)
		{
			DecisionTreeNode n = root.makeDecision(agent);
			return n.action;
		}
		
		public void printStructure()
		{
			root.print();
		}
		
		/* Private interface */
		
		/* returns the root of the tree */
		private void _makeTree(ArrayList examples, ArrayList attributes, DecisionTreeNode node) 
		{
			// Calculate the initial entropy
			float initialEntropy = this.entropy(examples);
			
			// if we have no entorpy, we can't divide further
			if (initialEntropy <= 0)
			{
				return;
			}
				
			// Find the number of examples
			float exampleCount = (float)examples.Count;
			
			// Hold the best found split so far
			float bestInformationGain = 0;
			Attribute bestSplitAttribute = Attribute.DEFAULT;
			
			Dictionary<bool, ArrayList> sets = new Dictionary<bool, ArrayList>();
			Dictionary<bool, ArrayList> bestSets = new Dictionary<bool, ArrayList>();
			
			// Go through each attribute
			foreach (Attribute attrib in attributes)
			{
				// Perform split	
				sets = splitByAttribute(examples, attrib);
				
				// Find the overall entropy and information gain
				float overallEntropy = this.entropyOfSets(sets, exampleCount);
				float informationGain = initialEntropy - overallEntropy;
				
				// Check if we've got the best so far
				if ( informationGain > bestInformationGain )
				{
					bestInformationGain = informationGain;
					bestSplitAttribute = attrib;
					bestSets = sets;
				}
			}
			
			// Set the decision node's test
			node.testAttribute = bestSplitAttribute;
			
			// The list of attributes to pass on down the tree shoud have the ones we're using removed
			ArrayList newAttributes = new ArrayList(attributes);
			newAttributes.Remove(bestSplitAttribute);
			
			// Fill the daughter nodes, based on the values in the "BestSets"
			foreach (KeyValuePair<bool, ArrayList> cur in bestSets)
			{
				// Find the value for the attribute in this set
				bool attributeValue = cur.Key;

				// if there is only one action left for this path, insert the action, not another decision
				Dictionary<Actions, int> actionTallies = this.actionsForExampleList(cur.Value);
				if (actionTallies.Count == 1)
				{
					//get the action
					foreach (KeyValuePair<Actions, int> curAt in actionTallies)
					{
						DecisionTreeNode daughter = new Action();
						daughter.nodeParent = node.nodeId;
						daughter.action = curAt.Key;
						
						// Add it to the tree - the action that should next b used if the agent has arived here
						node.daughters.Add(attributeValue,daughter);
					}
				}
				else
				{
					// Create a daughter node for the tree
					DecisionTreeNode daughter = new Decision();
					daughter.nodeParent = node.nodeId;
					// Add it to the tree - the nod that should next b visited if the agent has a given attributeValue
					node.daughters.Add(attributeValue,daughter);
			
					// Recurse the algorithm
					this._makeTree(cur.Value, newAttributes, daughter);					
				}
			}
		}
		
		/* splitByAttribute gives a dictionary with lists for each possible value
		 * The count of each list gives how many times that value occured
		 * splitByAttribute has been tested and works - Tom
		 * */
		private Dictionary<bool, ArrayList> splitByAttribute(ArrayList examples, Attribute attrib)
		{
			// We create a set of lists, so we can access each list
			// by the attribute value
			
			Dictionary<bool, ArrayList> sets = new Dictionary<bool, ArrayList>();
			
			// Loop through each example
			foreach (Example example in examples)
			{
				//check if there is a set for that example
				bool key = example.getValue(attrib);
				if(sets.ContainsKey(key))
				{
					//there is already an arraylist for this example so add to the example to that list
					sets[key].Add(example);
				}
				else
				{
					//this is the first example with this value for the attribute so make a new dict in sets
					ArrayList newList = new ArrayList();
					newList.Add(example);
					sets.Add(key, newList);
				}
			}
			
			// return the sets
			return sets;
		}
		
		public double log2(float x)
		{
			double res = Math.Log(x) / Math.Log(2.0);
			return res;			
		}
		
		private Dictionary<Actions, int> actionsForExampleList(ArrayList examples)
		{
			Dictionary<Actions, int> actionTallies = new Dictionary<Actions, int>();
			
			foreach (Example example in examples)
			{
				//check if there is a set for that example - wrong
				Actions key  = example.action;
				if(actionTallies.ContainsKey(key))
				{
					actionTallies[key] = actionTallies[key] + 1;
				}
				else
				{
					actionTallies.Add(key, 1);
				}
			}
			return actionTallies;		
		}
		
		/*
		 * entropy has been tested and works - Tom
		 * */
		private float entropy (ArrayList examples)
		{
			// Get the number of examples
			int exampleCount = examples.Count;
			
			// Check if we only have one
			if (exampleCount == 0)
			{
				return 0;
			}
			
			// Otherwise keep tally of how many action we have.
			Dictionary<Actions, int> actionTallies = this.actionsForExampleList(examples);
			int actionCount = actionTallies.Count;
			
			if (actionCount == 0)
			{
				return 0;
			}
			
			float entropyValue = 0;
			
			foreach (KeyValuePair<Actions, int> cur in actionTallies)
			{
				int tally = cur.Value;
				float proportion = tally/(float)exampleCount;
				float log2proportion = (float)log2(proportion);
				entropyValue -= proportion * log2proportion;
			}
			
			return entropyValue;
		}
		/*
		 * Tested and works - Tom
		 * 
		 * 
		 * */
		private float entropyOfSets(Dictionary<bool, ArrayList> sets, float exampleCount)
		{
			// Start with zero entropy
			float entropy = 0;
			
			// Get the entropy contribution of each set
			foreach (KeyValuePair<bool, ArrayList> cur in sets)
			{
				//Calculate the proportion of the whole in this set
				float proportion = cur.Value.Count/(float)exampleCount;
				//Calculate the entropy contribution
				entropy += proportion * this.entropy(cur.Value);
			}
			
			// return the total entropy
			return entropy;
		}
	}
}
